import json
import time
import os
from rag_chain import RAGChain
from text_processor import TextVectorizer

# 设置环境变量以避免认证问题
os.environ['TEST_JWT_TOKEN'] = 'test_token'

# 配置信息
TEST_PHONE = '13812345678'
EXPECTED_COMPANY = '上海科技发展有限公司'

# 主函数
def main():
    print(f'Final test for phone number {TEST_PHONE}')

    # 准备测试文档内容
    test_doc = "客户信息：公司名称上海科技发展有限公司，联系人张三，职位技术总监，电话号码13812345678，邮箱zhangsan@sh-tech.com。"
    print(f'Test document content: {test_doc}')

    # 初始化向量器和RAG链
    try:
        vectorizer = TextVectorizer()
        # 正确初始化RAGChain，传入api_key和vectorizer
        rag_chain = RAGChain(deepseek_api_key="fake_key", vectorizer=vectorizer)
        print('RAG chain initialized')
    except Exception as e:
        print(f'Error initializing RAG chain: {e}')
        return

    # 添加文档到向量器
    try:
        metadata = {'doc_id': 'test_doc', 'source': 'test_file.txt'}
        vectorizer.vectorize_and_store(test_doc, metadata)
        print('Document added to vectorizer')
    except Exception as e:
        print(f'Error adding document: {e}')
        return

    # 等待索引完成
    time.sleep(2)

    # 执行搜索
    try:
        print(f'Searching for phone number: {TEST_PHONE}')
        result = rag_chain.run(TEST_PHONE)
        print('Search completed')
    except Exception as e:
        print(f'Error during search: {e}')
        return

    # 打印完整结果
    print('\nFull search result:')
    print(json.dumps(result, ensure_ascii=False, indent=2))

    # 验证结果
    references = result.get('references', [])
    company_found = False
    similarity = 0

    print('\nVerification:')
    for ref in references:
        if EXPECTED_COMPANY in ref.get('full_text', ''):
            company_found = True
            similarity = ref.get('similarity', 0)
            print(f'SUCCESS: Found {EXPECTED_COMPANY} in references with similarity {similarity}')
            break

    if not company_found:
        print(f'FAILURE: Could not find {EXPECTED_COMPANY} in references')

if __name__ == '__main__':
    main()